from sklearn.metrics import roc_auc_score
import sys
sys.path.append("/home/cyw/projects/function_sim_project/TLSH/")
sys.path.append("/home/cyw/projects/function_sim_project/function_sim/")
sys.path.append("/home/cyw/projects/function_sim_project/Siamese_with_img/")
sys.path.append("/home/cyw/projects/function_sim_project/MGMN")
sys.path.append("/home/cyw/projects/function_sim_project/siamese_Graphsage")

from functionSim_train import functionModelWarpper
from hash_tlsh import mytlsh
from siamese_img_main import siamese_img_train
from MGMN_train import MGMNTrainer
from cfg_config import malSim_args
from siamese_graphsage_train import siamese_graphsage_trainer
from configure import *
# from functionsim_without_heterogeneous_train import function_without_heterogeous_train_Warpper
# from functionSim_with_cross_interaction_train import function_with_cross_interaction_train_Warpper

import shelve
from  easySample import easySample
import torch
import time
import copy
from tqdm import tqdm
from my_indicators import my_indicators

from show_experiment_res import img_tool


modelNames=["tlsh","functionSim","img","MGMN","siamese_graphsage","functionSim_with_hete","functionSim_zero","angr","retdec","radare2","ida"]
experIndic=["auc","precision","recall","mrr","ndcg","confusionMatrix","test_time","accuracy"]

class experimentRes():
    """
        用来生成对比试验数据的模型
    """

    def __init__(self,modelName,indicators,pairInf,pairName,indicatorsValue):
        """
            输入:
                加载最佳模型
                训练指标
                训练数据集+训练数据集名称
                训练指标参数
        """
        self.modelSample=easySample()
        self.modelName=modelName
        self.indicators=indicators
        self.myIndicators=my_indicators()
        self.indicatorsValue=indicatorsValue
        # 样本对的正确值
        self.values=[]
        for i in range(len(pairInf)):
            self.values.append(pairInf[i][2])
        self.pairInf=pairInf
        self.pairName=pairName
        self.model=None
        self.modelTrainer=None
        self.dataloader=None
        self.get_best_model(modelName)

    def get_best_model(self,modelName):
        """
            加载模型的最佳参数------>model
            部分模型需要加载Trainer--->modelTrain

            functionSim 的消融实验的config需要配置以下,不然这个不好直接修改!!!!!
        """
        if modelName=="tlsh":
            self.model=mytlsh()
        elif modelName=="functionSim":
            # self.model=torch.load(r'/home/cyw/projects/function_sim_project/all_data/models/functionSim_model.pth')
            self.model=torch.load(r'/home/cyw/projects/function_sim_project/all_data/models/functionSim_model_cross_edge_best_4_8.pth')
            self.modelTrainer=functionModelWarpper(self.model)
        elif modelName=="functionSim_with_hete":
            self.model=torch.load(r'/home/cyw/projects/function_sim_project/all_data/models/functionSim_with_hete_model.pth')
            self.modelTrainer=functionModelWarpper(self.model)
        elif modelName=="functionSim_zero":
            self.model=torch.load(r'/home/cyw/projects/function_sim_project/all_data/models/functionSim_zero_model.pth')
            self.modelTrainer=functionModelWarpper(self.model)
        elif modelName=="img":
            self.model=torch.load(r'/home/cyw/projects/function_sim_project/all_data/models/img_model.pth')
            # self.model=torch.load(r'/home/cyw/projects/function_sim_project/all_data/models/img_model_copy.pth')
            self.modelTrainer=siamese_img_train()
        elif modelName=="MGMN":
            self.model=torch.load(r'/home/cyw/projects/function_sim_project/all_data/models/MGMN_model.pth')
            self.modelTrainer=MGMNTrainer(self.model,malSim_args)
        elif modelName=="siamese_graphsage":
            self.model=torch.load(r'/home/cyw/projects/function_sim_project/all_data/models/siamese_graphsage_model.pth')
            self.modelTrainer=siamese_graphsage_trainer(self.model)
            self.dataloader=self.modelTrainer.get_model_dataloader(self.pairInf,self.pairName,False)
        elif modelName in ["angr","retdec","radare2","ida"]:
            self.model=torch.load(r'/home/cyw/projects/function_sim_project/all_data/models/{}_model.pth'.format(modelName))
            self.modelTrainer=functionModelWarpper(self.model)

    def get_res(self):
        """
            获得结果
            中转站
        """
        time1=time.time()
        if self.indicators=="auc":
           res=self.auc_res()
        elif self.indicators=="mrr":
            res=self.mrr_res(self.pairName,self.indicatorsValue)
        elif self.indicators=="ndcg":
            res=self.ndcg_res(self.indicatorsValue)
        elif self.indicators=="precision":
            res=self.precision_res_base_auc()
            # res=self.precision_res()
        elif self.indicators=="accuracy":
            res=self.accuracy_res_base_auc()
            # res=self.accuracy_res()
        elif self.indicators=="recall":
            res=self.recall_res(self.indicatorsValue)
        elif self.indicators=="confusionMatrix":
            res=self.confusion_matrix_res()
        elif self.indicators=="test_time":
            res=self.calculate_time_res(self.indicatorsValue)
        time2=time.time()
        with open("/home/cyw/projects/function_sim_project/all_logs/experiment_res_logs.txt","a") as file:
            file.write("        结果：{}    耗时：{}\n".format(res,time2-time1))
        print("     结果：{}    耗时：{}\n".format(round(res,5),round(time2-time1,5)))
    
    def get_model_calculate_time(self,pairInf,pairName,epochs):
        """
            用于获取各个模型的计算结果的耗时
        """
        print("为了测试结果的合理性，请移除模型中的批处理！！")
        if self.modelName=="tlsh":
            x,y=[],[]
            for j in range(epochs):
                for i in range(len(pairInf)):
                    x.append(self.modelSample.get_sample(pairInf[i][0],"tlsh"))
                    y.append(self.modelSample.get_sample(pairInf[i][1],"tlsh"))
            a=time.time()
            score=self.model.score(x,y)
        elif self.modelName=="img":
            imgDatas=self.modelTrainer.get_data(pairInf)
            img_batch_data=self.modelTrainer.trans_pair_into_batch(imgDatas)
            a=time.time()
            for i in range(epochs):
                es,temp=self.modelTrainer.get_scores_from_model(img_batch_data,self.model)
        elif self.modelName=="MGMN" or self.modelName=="siamese_graphsage":
            self.dataloader=self.modelTrainer.get_model_dataloader(pairInf,pairName,False)
            a=time.time()
            for i in range(epochs):
                res=self.modelTrainer.get_model_score(self.dataloader)
        else:
            self.dataloader=self.modelTrainer.get_model_dataloader(pairInf,pairName,False,False)
            a=time.time()
            for i in range(epochs):
                res=self.modelTrainer.get_model_score(self.dataloader)
        b=time.time()
        return b-a
    
    def get_model_res(self,pairInf,pairName):
        """
            用于获取各个模型的pairInf的输出结果
            确保一下，pairName并不会加载/
        """
        if self.modelName=="tlsh":
            x,y=[],[]
            for i in range(len(pairInf)):
                x.append(self.modelSample.get_sample(pairInf[i][0],"tlsh"))
                y.append(self.modelSample.get_sample(pairInf[i][1],"tlsh"))
            res=self.model.score(x,y)
            # res=self.myIndicators.modify_res(score)
        elif self.modelName=="img":
            imgDatas=self.modelTrainer.get_data(pairInf)
            img_batch_data=self.modelTrainer.trans_pair_into_batch(imgDatas)
            res,temp=self.modelTrainer.get_scores_from_model(img_batch_data,self.model)
        elif self.modelName=="MGMN" or self.modelName=="siamese_graphsage":
            self.dataloader=self.modelTrainer.get_model_dataloader(pairInf,pairName,False)
            res=self.modelTrainer.get_model_score(self.dataloader)
        else:
            self.dataloader=self.modelTrainer.get_model_dataloader(pairInf,pairName,False,False)
            res=self.modelTrainer.get_model_score(self.dataloader)
        return res

    def calculate_time_res(self,epochs):
        """
            模型平均计算一个样本的耗时
            modify-->移除数据加载时间
        """
        res=self.get_model_calculate_time(self.pairInf,self.pairName,epochs)
        ans=(res)*1.0/(len(self.pairInf)*epochs)
        print("总的训练耗时为{}".format(res))
        return ans

    def auc_res(self):
        """
            计算auc值
        """
        res=self.get_model_res(self.pairInf,self.pairName)
        aucValue=self.myIndicators.auc(self.values,res)    
        with shelve.open("/home/cyw/projects/function_sim_project/all_data/indicators/auc/{}_auc".format(self.modelName)) as file:
            file["pred"]=res
            file["true"]=self.values
        return aucValue
        
    # 下面三个函数,实现都比较取巧
    # 都直接使用MRR的计算结果
    def recall_res(self,recall_k):
        with shelve.open("/home/cyw/projects/function_sim_project/all_data/indicators/MRR/all_to_all_{}".format(self.modelName)) as file:
            ranked_list=file["ranked_list"]
            ground_truth=file["ground_truth"]
        res=0.0
        for i in range(len(ranked_list)):
            res+=self.myIndicators.recall_at_k(ranked_list[i],ground_truth[i],recall_k)
        ans=res/len(ranked_list)
        return ans


    def ndcg_res(self,ndcg_k):
        with shelve.open("/home/cyw/projects/function_sim_project/all_data/indicators/MRR/all_to_all_{}".format(self.modelName)) as file:
            ranked_list=file["ranked_list"]
            ground_truth=file["ground_truth"]
        res=0.0
        for i in range(len(ranked_list)):
            res+=self.myIndicators.ndcg_at_k(ranked_list[i],ground_truth[i],ndcg_k)
        ans=res/len(ranked_list)
        return ans
    
    def precision_res(self):
        with shelve.open("/home/cyw/projects/function_sim_project/all_data/indicators/MRR/all_to_all_{}".format(self.modelName)) as file:
            ranked_list=file["ranked_list"]
            ground_truth=file["ground_truth"]
        res=0.0
        for i in range(len(ranked_list)):
            res+=self.myIndicators.precision(ranked_list[i],ground_truth[i])
        ans=res/len(ranked_list)
        return ans
    
    def accuracy_res(self):
        with shelve.open("/home/cyw/projects/function_sim_project/all_data/indicators/MRR/all_to_all_{}".format(self.modelName)) as file:
            ranked_list=file["ranked_list"]
            ground_truth=file["ground_truth"]
        res=0.0
        for i in range(len(ranked_list)):
            res+=self.myIndicators.accuracy(ranked_list[i],ground_truth[i])
        ans=res/len(ranked_list)
        return ans
    
    def is_same_fam(self,x,y):
        """
            用于判断各个模型的分类结果是否属于同一个家族
        """
        limit_num=0.8
        VALUE_TEMP=0.5
        if self.modelName=="tlsh":
            if (x>=0.5 and y>VALUE_TEMP)  or (x<0.5 and y<VALUE_TEMP):
                return True
        elif self.modelName=="img":
            if (x>=limit_num and y>VALUE_TEMP)  or (x<limit_num and y<VALUE_TEMP):
                return True
        elif self.modelName=="MGMN":
            if (x>=limit_num and y>VALUE_TEMP)  or (x < limit_num and y<VALUE_TEMP):
                return True
        elif self.modelName in ["functionSim","functionSim_zero","functionSim_with_hete"]:
            if (x>=limit_num and y>VALUE_TEMP)  or (x<limit_num and y<VALUE_TEMP):
                return True
        elif self.modelName=="siamese_graphsage":
            if (x>=limit_num and y>VALUE_TEMP)  or (x<limit_num and y<VALUE_TEMP):
                return True
        return False
    
    def accuracy_res_base_auc(self):
        res=self.get_model_res(self.pairInf,self.pairName)
        size=len(self.values)
        tp,tn,fp,fn=0,0,0,0
        for i in range(size):
            tar=self.values[i]
            temp=res[i]
            if self.is_same_fam(temp,tar):
                if tar>0.5:
                    tp+=1
                else:
                    tn+=1
            else:
                if tar<0.5:
                    fn+=1
                else:
                    fp+=1
        return (tp+tn)*1.0/(tp+fp+tn+fn)
    
    def precision_res_base_auc(self):
        res=self.get_model_res(self.pairInf,self.pairName)
        size=len(self.values)
        tp,tn,fp,fn=0,0,0,0
        for i in range(size):
            tar=self.values[i]
            temp=res[i]
            if self.is_same_fam(temp,tar):
                if tar>0.5:
                    tp+=1
                else:
                    tn+=1
            else:
                if tar<0.5:
                    fn+=1
                else:
                    fp+=1
        return tp*1.0/(tp+fn)
    
    
    def confusion_matrix_res(self):
        with shelve.open("/home/cyw/projects/function_sim_project/all_data/indicators/MRR/all_to_all_{}".format(self.modelName)) as file:
            ranked_list=file["ranked_list"]
            ground_truth=file["ground_truth"]
            sim_res=file["sim_res"]
            name_list=file["nameList"]
            name_to_lable=file["name_to_lable"]
            lable_to_name=file["lable_to_name"]

        # 是乱序的，这里按照家族来合并一下
        lth=len(lable_to_name)#家族的个数
        lable_to_ind={}
        name_to_ind={}
        cnt=0
        lable_nums={}
        lables=[]
        for i in lable_to_name.keys():
            lables.append(i)
            lable_to_ind[i]=cnt
            cnt+=1
            lable_nums[i]=len(lable_to_name[i])*(len(lable_to_name[i])-1)
        cnt=0
        for i in range(len(name_list)):
            name_to_ind[name_list[i]]=cnt#这份是啥，我忘记了
            cnt+=1

        # 计算混淆矩阵    
        # 同时保存误报、漏报信息
        loubao={}
        matrix_res=[[0] * (lth+1) for i in range(lth+1)]
        for family in lables:
            a=lable_to_ind[family]
            for name in lable_to_name[family]:
                ind=name_to_ind[name]
                x=ranked_list[ind]
                size=len(ground_truth[ind])
                for i in range(size):
                    tar=x[i]
                    b=lable_to_ind[name_to_lable[tar]]
                    matrix_res[a+1][b+1]+=1
                    if a!=b:
                        if family not in loubao:
                            loubao[family]=[(name,tar,name_to_lable[tar])]
                        else:
                            loubao[family].append((name,tar,name_to_lable[tar]))
        for i in range(1,lth+1):
            matrix_res[0][i]=lables[i-1]
            matrix_res[i][0]=lables[i-1]
        print("正在绘制表格")
        gene_img=img_tool()
        gene_img.draw_confusion_matrix(self.modelName,matrix_res,lables)
        print("表格绘制完成")
        with open("/home/cyw/projects/function_sim_project/all_data/indicators/loubao/loubao.txt","a") as file:
            file.write("{}  的漏报信息：\n".format(self.modelName))
            for i in loubao.keys():
                temp={}
                for originName,name,fam in loubao[i]:
                    if fam not in temp:
                        temp[fam]=[originName+"--"+name]
                    else:
                        temp[fam].append(originName+"--"+name)
                file.write("    家族名：{}\n".format(i))
                for j in temp.keys():
                    file.write("        误判家族名：{}\n".format(j))
                    res=" ".join(temp[j])
                    file.write("            {}\n".format(res))
        print("数据保存成功")
        return True

 
    def mrr_res(self,saveName,use_save_data):
        """
            一个样本和所有的样本进行相似性的度量,然后保留对与对之间的相似值
            计算出每个家族所包含的所有样本,做成list,当成ground Truth

            use_save_data,表示是否使用之前保存的数据
            
        """
        if use_save_data=="use_save_data":
            with shelve.open("/home/cyw/projects/function_sim_project/all_data/indicators/MRR/all_to_all_{}".format(self.modelName)) as file:
                ranked_list=file["ranked_list"]
                ground_truth=file["ground_truth"]
        else:
            if gene_disassemble_pair == False:
                with shelve.open(r"/home/cyw/projects/function_sim_project/all_data/pair_infs/{}".format(saveName)) as file:
                    allPairs=file["pairs"]
                    lable_to_name=file["lable_to_name"]
                    name_to_lable=file["name_to_lable"]
                    ans=[]
            else:
                with shelve.open(r"/home/cyw/projects/function_sim_project/all_data/pair_infs/dis_{}".format(saveName)) as file:
                    allPairs=file["pairs"]
                    lable_to_name=file["lable_to_name"]
                    name_to_lable=file["name_to_lable"]
                    ans=[]

            # 每一个pairs就是一个样本和其他所有样本的组合
            for pairs in tqdm(allPairs):
                temp=[]
                # 这个应该是test
                res=self.get_model_res(pairs,self.pairName)
                for i in range(len(res)):
                    temp.append((res[i],pairs[i][0],pairs[i][1]))
                # 按照相似值的大小排序
                temp.sort(reverse=True)
                ans.append(copy.deepcopy(temp))
            
            # 计算mrr值
            ranked_list=[]
            ground_truth=[]
            nameList=[]
            for i in range(len(ans)):
                temp=[]
                for j in range(len(ans[i])):
                    temp.append(ans[i][j][2])
                ranked_list.append(temp)
                ground_truth.append(lable_to_name[name_to_lable[ans[i][0][1]]])
                nameList.append(ans[i][0][1])

            #   保存计算数据
            with shelve.open("/home/cyw/projects/function_sim_project/all_data/indicators/MRR/all_to_all_{}".format(self.modelName)) as file:
                file["ranked_list"]=ranked_list
                file["ground_truth"]=ground_truth
                file["nameList"]=nameList
                # res没用。。。这里只保存了最后一行的相似值
                file["sim_res"]=res
                file["lable_to_name"]=lable_to_name
                file["name_to_lable"]=name_to_lable

        return self.myIndicators.MRR(ranked_list,ground_truth)

def getPairInf(dataName,path):
    with shelve.open(path) as file:
        if dataName=="valid":
            testPair=file["validPair"]
        elif dataName=="test":
            testPair=file["testPair"]
        # 移除家族信息
        pairInf = [row[:3] for row in testPair]
    return pairInf
def run_for_res(modelNames,indics):
    """ 
        用来获得各个模型的各个指标值

        输入:[模型名],[指标名],以及对应指标的参数
        输出:无,保存到日志文件中
    """
    dataName="test"
    # dataName="valid"
    
    # 获得用于生成指标的数据对
    if gene_disassemble_pair == False:
        pairInf = getPairInf(dataName,"/home/cyw/projects/function_sim_project/all_data/pair_infs/pair_infs")
    else:
        # 加载反汇编实验所需的样本对
        pairInf = getPairInf(dataName,"/home/cyw/projects/function_sim_project/all_data/pair_infs/dis_pair_infs")

    for modelName in modelNames:
        if modelName in ["functionSim_zero","functionSim_with_hete"]:
            print("functionSim消融实验，注意配置信息是否切换！！！")
        for indic in indics.keys():
            # 部分指标有参数需要设置
            value=indics[indic]
            with open("/home/cyw/projects/function_sim_project/all_logs/experiment_res_logs.txt","a") as file:
                file.write("{}模型  计算{}值:\n".format(modelName,indic))
            print("{}模型  计算{}值:\n".format(modelName,indic))
            tools=experimentRes(modelName,indic,pairInf,dataName,value)
            tools.get_res()

# 运行命令:nohup python /home/cyw/projects/function_sim_project/basic_script/experiment_res.py
# 消融实验部分代码逻辑未完成,模型指标能放跑通未验证
if __name__=="__main__":
    """
        评价指标:
            auc         √
            precision   √
            recall@k    √
            mrr         √
            ndcg@k      √
            testTime    √（感觉不太正确，部分模型样本加载太慢了）        
    """

    # modelName=["siamese_graphsage"]
    # modelName=["functionSim_zero"]
    # modelName=["functionSim_with_hete"]
    modelName=["functionSim"]
    # /home/cyw/projects/function_sim_project/all_data/models/functionSim_model_cross_edge.pth
    # modelName=["MGMN"]
    # modelName=["img"]
    # modelName=["tlsh"]
    # modelName=["tlsh","img","MGMN","functionSim","siamese_graphsage"]
    # modelName=["functionSim","siamese_graphsage"]
    # modelName = ["angr","retdec","radare2","ida"]
    # modelName = ["angr"]
    # modelName = ["retdec"]
    # modelName = ["ida"]
    # modelName = ["radare2"]
    indic={
        # "auc":True,
        # "mrr":"else",
        "mrr":"use_save_data",
        # "precision":True,
        # "accuracy":True,
        "ndcg":10,
        "recall":10,
        # "test_time":10
    }
    is_able_run=True
    for name in modelName:
        if name not in modelNames:
            is_able_run=False
    for ind in indic.keys():
        if ind not in experIndic:
            is_able_run=False
    if is_able_run:
        run_for_res(modelName,indic)
    else:
        print("模型或者指标指定错误！！！")